# -*- coding: utf-8 -*-
# ===========================================
# @Time    : 2021/8/24 上午11:57
# @Author  : shutao
# @FileName: data_augment.py
# @remark  : 
# 
# @Software: PyCharm
# Github 　： https://github.com/NameLacker
# ===========================================

from loguru import logger
import numpy as np
import cv2 as cv
import random
import json
import os


def _rotate(points, m):
    """
    坐标点映射
    :param points:
    :param m:
    :return:
    """
    out_points = []
    for ps in points:
        pts = np.float32(ps).reshape([-1, 2])  # 要映射的点
        pts = np.hstack([pts, np.ones([len(pts), 1])]).T
        target_point = np.dot(m, pts)
        target_point = [[target_point[0][x], target_point[1][x]] for x in range(len(target_point[0]))]
        out_points.append(target_point)
    return out_points


def _rotate_bound(image, angle):
    """
    旋转且不裁剪原图
    :param image:
    :param angle:
    :return:
    """
    # grab the dimensions of the image and then determine the
    # center
    (h, w) = image.shape[:2]
    (cX, cY) = (w // 2, h // 2)

    # grab the rotation matrix (applying the negative of the
    # angle to rotate clockwise), then grab the sine and cosine
    # (i.e., the rotation components of the matrix)
    M = cv.getRotationMatrix2D((cX, cY), -angle, 1.0)
    cos = np.abs(M[0, 0])
    sin = np.abs(M[0, 1])

    # compute the new bounding dimensions of the image
    nW = int((h * sin) + (w * cos))
    nH = int((h * cos) + (w * sin))

    # adjust the rotation matrix to take into account translation
    M[0, 2] += (nW / 2) - cX
    M[1, 2] += (nH / 2) - cY

    # perform the actual rotation and return the image
    return cv.warpAffine(image, M, (nW, nH)), M


def _img_whirl(img_file, lab_file, angle=random.randint(10, 350), resize_rate=1):
    """
    图像旋转
    :return:
    """
    # 读取坐标点列表
    with open(lab_file, 'r') as f:
        lab = json.load(f)
    seg_labels = lab["shapes"]
    points = [seg_label["points"] for seg_label in seg_labels]

    # 读取图像
    img = cv.imread(img_file)
    # 旋转
    res_img, M = _rotate_bound(img, angle)

    # 坐标点映射回对应标签
    points = _rotate(points, M)
    for idx, seg_label in enumerate(lab["shapes"]):
        seg_label["points"] = points[idx]
    return res_img, lab


def image_augment(images_path="E:\\work\\Dataset\\house_dataset\\images",
                  labels_path="E:\\work\\Dataset\\house_dataset\\seg_images",
                  NOT=5):
    """
    数据增广
    Args:
        images_path: 图像目录
        labels_path: 标签目录
        NOT: 增广次数

    Returns:

    """
    logger.info("开始增广数据量...")
    for NoT in range(NOT):
        for idx, _file in enumerate(os.listdir(images_path)):
            if 'rotate' in _file:
                continue
            logger.info(_file)
            image_file = os.path.join(images_path, _file)
            label_file = os.path.join(labels_path, _file.split('.')[0] + ".json")

            res_img, label = _img_whirl(image_file, label_file)
            cv.imwrite(os.path.join(images_path, 'rotate_{}_{}'.format(NoT+1, _file)), res_img)
            with open(os.path.join(labels_path, "rotate_{}_{}.json".format(NoT+1, _file.split('.')[0])), 'w') as f:
                json.dump(label, f, indent=4)
